<?php
/**
 * Message queue consumer, executing image generation commands.
 *
 * Consumers will execute commands defined in the configuration file and pass
 * parameters passed from the front-end to the back-end. Before and after the
 * command execution, the information of the current consumption queue will be
 * obtained and written into the directory specified in the configuration file.
 */
namespace App\Jobs;

use Throwable;
use Illuminate\Support\Facades\Log;
use Illuminate\Support\Facades\Config;
use Illuminate\Bus\Queueable;
use Illuminate\Contracts\Queue\ShouldQueue;
use Illuminate\Foundation\Bus\Dispatchable;
use Illuminate\Queue\InteractsWithQueue;
use Illuminate\Queue\SerializesModels;

class ProcessImageGenerate implements ShouldQueue
{
    use Dispatchable, InteractsWithQueue, Queueable, SerializesModels;

    private string $fileName;
    private string $prompt;
    private int $width;
    private int $height;
    private int $steps;
    private int $seed;
    private int $norwkv;
    private string $sdExe;
    private string $sdParams;
    private string $writeFloder;

    /**
     * Create a new job instance.
     *
     * @param string $name     The file name, usually in the format of
     *                         "20240501-142001-19798332"
     * @param string $prompt   Prompt words submitted from the front-end
     * @param int    $width    The width of the image should be a multiple of 64
     * @param int    $height   The height of the image should be a multiple of 64
     * @param int    $steps    The number of sampling steps should be greater than 0
     * @param int    $seed     The seed of a random number, which should be less
     *                         than 2^32-1
     * @param int    $norwkv   Is RWKV used to process prompt words? If it is 0,
     *                         then it is not used; otherwise, it is used.
     * @return void
     */
    public function __construct(
        string $name,
        string $prompt = '', int $width = 512, int $height = 512,
        int $steps = 1, int $seed = -1,
        int $norwkv = 0
    ) {
        $this->fileName = $name;
        $this->prompt = $prompt;
        $this->width = $width;
        $this->height = $height;
        $this->steps = $steps;
        $this->seed = $seed;
        $this->norwkv = $norwkv;
        $this->sdExe = trim(Config::get('images.sd_exe'));
        $this->sdParams = trim(Config::get('images.sd_params'));
        $this->writeFloder = trim(Config::get('images.dir'));
        Log::info('ProcessImageGenerate: '.json_encode([
            'fileName' => $this->fileName,
            'prompt' => $this->prompt,
            'width' => $this->width,
            'height' => $this->height,
            'steps' => $this->steps,
            'seed' => $this->seed,
            'norwkv' => $this->norwkv,
            'sdExe' => $this->sdExe,
            'sdParams' => $this->sdParams,
            'writeFloder' => $this->writeFloder,
        ]));
    }

    /**
     * Execute the job.
     *
     * @return void
     */
    public function handle()
    {
        $this->generateImage();
    }

    /**
     * Execute image generation commands
     *
     * @return void
     */
    private function generateImage()
    {
        // It is necessary to ensure that the parameter of the constructor
        // is after the parameters in the user configuration file. Because it
        // allows users to modify configuration files and execute other scripts,
        // such as Python scripts.
        $command = array_merge(
            [$this->sdExe],
            explode(' ', $this->sdParams),
            [
                '-o', $this->writeFloder.DIRECTORY_SEPARATOR.$this->fileName.'.png',
                '-p', $this->prompt,
                '--steps', $this->steps,
                '-s', $this->seed,
                '-W', $this->width,
                '-H', $this->height,
                '--norwkv', $this->norwkv,
            ]
        );
        // I don't know why, it must be written to a file,
        // otherwise the process will stop running when paraphrasing with rwkv.
        $commandOutputTmpFile = tempnam(sys_get_temp_dir(), 'stdout');
        $commandErrorOutputTmpFile = tempnam(sys_get_temp_dir(), 'stderr');

        if (false === ($processResource = proc_open(
            $command,
            [['pipe', 'r'], ['file', $commandOutputTmpFile, 'wb'], ['file', $commandErrorOutputTmpFile, 'wb']],
            $pipes
        ))) {
            Log::error('Failed to execute image generation task, process unable to execute: '
            . json_encode($command));
            return;
        }

        $status = ['running' => true, 'exitcode' => -1];
        do {
            if ($status['running']) {
                sleep(1);
            }
            $status = proc_get_status($processResource);
        } while ($status['running']);

        if (0 != $status['exitcode']) {
            $info = file_get_contents($commandOutputTmpFile);
            $error = file_get_contents($commandErrorOutputTmpFile);
            Log::error('Failed to execute image generation task, The process '
            . 'has completed execution but the exit code = '
            . $status['exitcode'] . '. Command = ' . json_encode($command)
            . ' stdout = ' . json_encode($info)
            . ' errout = ' . json_encode($error));
            @proc_close($processResource);
            @unlink($commandOutputTmpFile);
            @unlink($commandErrorOutputTmpFile);
            return;
        }

        @proc_close($processResource);
        @unlink($commandOutputTmpFile);
        @unlink($commandErrorOutputTmpFile);
        file_put_contents(
            $this->writeFloder.DIRECTORY_SEPARATOR.$this->fileName.'.json',
            json_encode($command)
        );
    }
}
